"""
    给基因簇识别和大模型总结表格添加概念树注释的6级分类
"""


import json
import argparse
import pandas as pd


def label_to_tree_concept(add_pfam_table, cluster_tble, concept_tree, output):
    with open(concept_tree) as file:
        label_tree = json.load(file)

    # for k,v in label_tree.items():
    #     print(k,v)

    cluster_df = pd.read_excel(cluster_tble)
    cluster_dict = {}
    for row in cluster_df.values.tolist():
        # print(row)
        row_3 = eval(
            "['" +
            row[2].strip('"]').strip("'").strip("['").strip('"').replace("', '",
                "---").replace('"', "'").replace("'", "").replace("---", "', '")
            + "']"
        )
        row_1 = row[0]
        tree_dict = {
            1: [],
            2: [],
            3: [],
            4: [],
            5: [],
            6: []
        }
        for l in row_3:
            l = l.strip()
            if l in label_tree:
                tree_concepts = label_tree[l]["concepts"]
                if isinstance(tree_concepts, dict):
                    tree_dict[1].append(tree_concepts["concept"])
                else:
                    for index, concept in enumerate(tree_concepts):
                        tree_dict[index + 1].append(concept)

        tree_list = [""]*6
        for k,v in tree_dict.items():
            tree_list[k-1] = "||".join(list(set(v)))
        #print(tree_list)
        cluster_dict[row_1] = row[1:] + tree_list

    all_df = pd.read_excel(add_pfam_table)
    title = all_df.columns.tolist()

    result = []
    for row in all_df.values.tolist():
        cluster = row[5]
        #cluster_tree = cluster_dict[cluster]
        # 文件中是"cluster_1"，需要在前面加上，不然匹配不上
        cluster_tree = cluster_dict["cluster_"+str(cluster)]
        row = row + cluster_tree
        result.append(row)

    title = title + ["main_category","summary_feature","level_1",
                    "level_2","level_3","level_4","level_5","level_6"]
    result_df = pd.DataFrame(result, columns=title)
    result_df.to_excel(output, index=None)
    print(output)


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("-add_pfam_table", "--add_pfam_table")
    parser.add_argument("-cluster_table", "--cluster_table")
    parser.add_argument("-concept_tree", "--concept_tree")
    parser.add_argument("-output","--output")
    params = parser.parse_args()

    return params


if __name__ == "__main__":
    params = parse_args()
    add_pfam_table = params.add_pfam_table
    cluster_tble = params.cluster_table
    concept_tree = params.concept_tree
    output = params.output

    label_to_tree_concept(add_pfam_table, cluster_tble, concept_tree, output)

